package com.langChain4j.rag.config;

import com.langChain4j.api.constants.ModelHostConstant;
import com.langChain4j.api.enums.ModelEnum;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * 向量数据存储方式配置
 * @author linjun
 * @date 2025/9/25 14:24
 * @description
 */
@Configuration
public class EmbeddingStoreConfiguration {

    @Value("${langchain4j.embeddingStore.chroma.host:http://localhost:18000}")
    private String chromaHost;

    @Autowired
    private RedisProperties redisProperties;

    @Autowired
    private PgVectorProperties pgVectorProperties;

    @Bean
    @ConditionalOnProperty(name = "langchain4j.embeddingStore.redis.enabled", havingValue = "true")
    public EmbeddingStore<TextSegment> redisEmbeddingStore() {
        return RedisEmbeddingStore.builder()
                .host(redisProperties.getHost())
                .port(redisProperties.getPort())
                .password(redisProperties.getPassword())
                .dimension(384)
                .build();
    }

    @Bean
    @ConditionalOnProperty(name = "langchain4j.embeddingStore.chroma.enabled", havingValue = "true")
    public EmbeddingStore<TextSegment> chromaEmbeddingStore() {
        return ChromaEmbeddingStore.builder()
                .baseUrl(chromaHost)
                .collectionName("chromaCollection")
                .logRequests(true)
                .logResponses(true)
                .build();
    }

//    @Bean
//    @ConditionalOnProperty(name = "langchain4j.embeddingStore.memory.enabled", havingValue = "true", matchIfMissing = true)
//    @ConditionalOnMissingBean(EmbeddingStore.class)
//    public EmbeddingStore<TextSegment> memoryEmbeddingStore() {
//        return new InMemoryEmbeddingStore<>();
//    }

    /**
     * pgVector存储模型
     * @return
     */
    @Bean
    @ConditionalOnProperty(name = "langchain4j.embeddingStore.pgvector.enabled", havingValue = "true")
    @ConditionalOnMissingBean(EmbeddingStore.class)
    public EmbeddingStore<TextSegment> pgVectorEmbeddingStore() {
        return PgVectorEmbeddingStore.builder()
                .host(pgVectorProperties.getHost())
                .port(pgVectorProperties.getPort())
                .database(pgVectorProperties.getDatabase())
                .user(pgVectorProperties.getUser())
                .password(pgVectorProperties.getPassword())
                .table("embedding")
                .dimension(pgVectorProperties.getDimension())
                .dropTableFirst(true)
                .createTable(true)
                .build();
    }

    /**
     * 向量模型
     * @return
     */
    @Bean
    public EmbeddingModel embeddingModel() {
        OllamaEmbeddingModel.OllamaEmbeddingModelBuilder ollamaEmbeddingModelBuilder = OllamaEmbeddingModel.builder()
                .modelName(ModelEnum.EMBEDDING_MODEL_OLLAMA_MINI_LM.getCode())
                .logRequests(true)
                .logResponses(true)
                .baseUrl(ModelHostConstant.OLLAMA_API_URL);
        return ollamaEmbeddingModelBuilder.build();
    }

}
